import os
import requests
import m3u8
from Crypto.Cipher import AES
from concurrent.futures import ThreadPoolExecutor
import subprocess
import shutil

# 配置
# M3U8_URL = "https://v30.51cto.com/c0e39f05fbe271ef80335017e1e90102/fb3ce53c9d9009c919bd8dccc63e837c-hd-encrypt-stream.m3u8"
M3U8_URL = "https://vod1.xisaiwang.cn/202505252051/c6afd93ea4e210964971d6397db6069a/e01ed63cf29171ef801916b5feac0102/237404f91d07e082c591b90da4c7aebb-S00000001-100000-od-encrypt-stream.m3u8"

SAVE_DIR = "./downloaded_video"
THREADS = 10  # 下载线程数

# 创建主保存目录
os.makedirs(SAVE_DIR, exist_ok=True)

# TS 分片专用目录
TS_SEGMENTS_DIR = os.path.join(SAVE_DIR, "ts_segments")
os.makedirs(TS_SEGMENTS_DIR, exist_ok=True)

# 1. 解析 M3U8 文件
def parse_m3u8():
    playlist = m3u8.load(M3U8_URL)
    key = None
    if playlist.keys and playlist.keys[0]:
        key_url = playlist.keys[0].uri
        base_url = M3U8_URL.rsplit("/", 1)[0]
        key_url = base_url + "/" + key_url if not key_url.startswith("http") else key_url
        key_response = requests.get(key_url)
        key = (key_url, key_response.content)  # (key_url, key_bytes)
    ts_list = [seg.absolute_uri for seg in playlist.segments]
    return key, ts_list

# 2. 下载并解密 TS 分片
def download_ts(ts_url, index, key, iv=None):
    try:
        headers = {"User-Agent": "Mozilla/5.0"}
        response = requests.get(ts_url, headers=headers, stream=True)
        if response.status_code != 200:
            print(f"下载失败 {ts_url}")
            return

        data = response.content
        if key:
            cipher = AES.new(key[1], AES.MODE_CBC, iv=key[1] if iv is None else iv)
            data = cipher.decrypt(data)
        filename = os.path.join(TS_SEGMENTS_DIR, f"segment_{index:04d}.ts")
        with open(filename, "wb") as f:
            f.write(data)
        print(f"下载完成 {filename}")
    except Exception as e:
        print(f"下载错误 {ts_url}: {e}")

# 3. 合并 TS 分片
def merge_ts_files(output_mp4):
    ts_files = sorted([os.path.join(TS_SEGMENTS_DIR, f) for f in os.listdir(TS_SEGMENTS_DIR) if f.endswith(".ts")])
    filelist_path = os.path.join(TS_SEGMENTS_DIR, "filelist.txt")

    with open(filelist_path, "w") as f:
        for ts in ts_files:
            f.write(f"file '{os.path.basename(ts)}'\n")

    command = [
        "ffmpeg",
        "-f", "concat",
        "-safe", "0",
        "-i", filelist_path,
        "-c", "copy",
        output_mp4
    ]
    try:
        subprocess.run(command, check=True)
        print(f"合并完成: {output_mp4}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"FFmpeg 合并失败: {e}")
        return False

# 4. 主流程
def main():
    key, ts_list = parse_m3u8()
    print(f"共找到 {len(ts_list)} 个 TS 分片")

    # 多线程下载
    with ThreadPoolExecutor(max_workers=THREADS) as executor:
        for i, ts_url in enumerate(ts_list):
            executor.submit(download_ts, ts_url, i, key)

    output_mp4 = os.path.join(SAVE_DIR, "output.mp4")

    # 合并并清理
    if merge_ts_files(output_mp4):
        try:
            shutil.rmtree(TS_SEGMENTS_DIR)
            print("TS 分片目录已删除")
        except Exception as e:
            print(f"删除 TS 目录失败: {e}")

if __name__ == "__main__":
    main()